/**
 *
 */
package org.windowkit.toolkit.widget.tree;

import java.awt.Color;
import java.awt.Component;
import java.awt.Dimension;
import java.awt.Graphics;
import java.awt.Rectangle;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.Hashtable;
import java.util.Vector;

import javax.swing.ButtonGroup;
import javax.swing.Icon;
import javax.swing.JCheckBoxMenuItem;
import javax.swing.JMenu;
import javax.swing.JMenuItem;
import javax.swing.JPopupMenu;
import javax.swing.JTree;
import javax.swing.plaf.TreeUI;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.TreeModel;
import javax.swing.tree.TreeNode;
import javax.swing.tree.TreePath;
import javax.swing.tree.TreeSelectionModel;

/**
 * @author Rich
 *
 */
public class CheckBoxTree extends JTree {

	/**
	 *
	 */
	private static final long serialVersionUID = -7180799588022336076L;

	CheckBoxTreeCellRenderer cellRenderer;

	Vector<DefaultMutableTreeNode> checkedNodes;

    protected JPopupMenu popupMenu;

    private final DefaultTreeModel EMPTY_TREE =
        new DefaultTreeModel(new DefaultMutableTreeNode());



	/**
	 *
	 */
	public CheckBoxTree() {
		this(new Object[]{""});
	}

	/**
	 * @param value
	 */
	public CheckBoxTree(final Object[] value) {
		super(value);
		init();
	}

	/**
	 * @param value
	 */
	public CheckBoxTree(final Vector<?> value) {
		super(value);
		init();
	}

	/**
	 * @param value
	 */
	public CheckBoxTree(final Hashtable<?, ?> value) {
		super(value);
		init();
	}

	/**
	 * @param root
	 */
	public CheckBoxTree(final TreeNode root) {
		super(root);
		init();
	}

	/**
	 * @param newModel
	 */
	public CheckBoxTree(final TreeModel newModel) {
		super(newModel);
		init();
	}

	/**
	 * @param root
	 * @param asksAllowsChildren
	 */
	public CheckBoxTree(final TreeNode root, final boolean asksAllowsChildren) {
		super(root, asksAllowsChildren);
		init();
	}

	private void init(){
		checkedNodes = new Vector<DefaultMutableTreeNode>();
		cellRenderer = new CheckBoxTreeCellRenderer();
		setCellRenderer(cellRenderer);
        getSelectionModel().setSelectionMode(TreeSelectionModel.SINGLE_TREE_SELECTION);

		addMouseListener(new CheckBoxTreeListener());

        setUpPopUp();
     }

    protected void setUpPopUp(){
        // set up the popup menu
        popupMenu = new JPopupMenu();
        JMenuItem item = new JMenuItem("Expand All");
        popupMenu.add(item);
        item.addActionListener(new ActionListener(){
            public void actionPerformed(final ActionEvent e) {
                expandAll();
            }
        });

        item = new JMenuItem("Collapse All");
        popupMenu.add(item);
        item.addActionListener(new ActionListener(){
           public void actionPerformed(final ActionEvent e) {
               collapseAll();
           }
        });

        popupMenu.addSeparator();

        item = new JMenuItem("Check All");
        item.addActionListener(new ActionListener(){
           public void actionPerformed(final ActionEvent e) {
               checkAllNodes();
           }
        });
        popupMenu.add(item);

        item = new JMenuItem("Uncheck All");
        item.addActionListener(new ActionListener(){
            public void actionPerformed(final ActionEvent e) {
                uncheckAllNodes();
            }
        });
        popupMenu.add(item);

        popupMenu.addSeparator();

        popupMenu.add(createColorItem());


        //add the popup menu listener
        addMouseListener(new MouseAdapter(){
            @Override
			public void mouseReleased(final MouseEvent e) {
                if(e.isPopupTrigger()) {
					popupMenu.show(e.getComponent(), e.getX(), e.getY());
				}
            }
        });


    }

    private JMenu createColorItem(){
        JMenu colorMenu = new JMenu("Color");

        ButtonGroup grp = new ButtonGroup();

        JCheckBoxMenuItem item = new JCheckBoxMenuItem("Light Gray", new ColorIcon(Color.LIGHT_GRAY));
        item.addActionListener(new ColorChanger(Color.LIGHT_GRAY));
        //this is the starting color
        item.setState(true);
        colorMenu.add(item);
        grp.add(item);


        item = new JCheckBoxMenuItem("Blue", new ColorIcon(Color.BLUE));
        item.addActionListener(new ColorChanger(Color.BLUE));
        colorMenu.add(item);
        grp.add(item);

        item = new JCheckBoxMenuItem("Blood", new ColorIcon(Color.RED));
        item.addActionListener(new ColorChanger(Color.RED));
        colorMenu.add(item);
        grp.add(item);

        item = new JCheckBoxMenuItem("Black", new ColorIcon(Color.BLACK));
        item.addActionListener(new ColorChanger(Color.BLACK));
        colorMenu.add(item);
        grp.add(item);

        item = new JCheckBoxMenuItem("Cyan", new ColorIcon(Color.CYAN));
        item.addActionListener(new ColorChanger(Color.CYAN));
        colorMenu.add(item);
        grp.add(item);

        item = new JCheckBoxMenuItem("Ecto", new ColorIcon(Color.GREEN));
        item.addActionListener(new ColorChanger(Color.GREEN));
        colorMenu.add(item);
        grp.add(item);



        item = new JCheckBoxMenuItem("Golden Harvest", new ColorIcon(Color.ORANGE));
        item.addActionListener(new ColorChanger(Color.ORANGE));
        colorMenu.add(item);
        grp.add(item);


        return colorMenu;
    }



     public void expandAll(){
         for(int i = 0; i < getRowCount(); i++) {
			expandRow(i);
		}
     }

     public void collapseAll(){
         TreePath[] paths = getPathBetweenRows(0, getRowCount());
         for(int i = paths.length - 1; i >= 0; i--) {
			collapsePath(paths[i]);
		}

     }

     public void checkAllNodes(){
         DefaultMutableTreeNode root = (DefaultMutableTreeNode)getModel().getRoot();
         checkedNodes.clear();
         if(root != null){
             checkNode(root);
         }

     }

     public void uncheckAllNodes(){
         checkedNodes.clear();
         repaint();
     }

     public void eraseTree(){
         setModel(EMPTY_TREE);
     }





	void checkNode(final DefaultMutableTreeNode node){
		if(!checkedNodes.contains(node)){
            //check
			checkedNodes.add(node);
            if(!node.isLeaf()){
                expandRow(getRowForPath(new TreePath(node.getPath())));
                for(Enumeration<DefaultMutableTreeNode> children = node.children();
                children.hasMoreElements(); ){
                    DefaultMutableTreeNode child = children.nextElement();
                    if(!isNodeChecked(child)){
                        checkNode(child);
                    }
                }
            }
		}else{
            //uncheck
			checkedNodes.remove(node);
            if(!node.isLeaf()){
                for(Enumeration<DefaultMutableTreeNode> children = node.children();
                children.hasMoreElements(); ){
                    DefaultMutableTreeNode child = children.nextElement();
                    if(isNodeChecked(child)){
                        checkNode(child);
                    }
                }
            }
            DefaultMutableTreeNode parent =(DefaultMutableTreeNode)node.getParent();
            if(parent != null && isNodeChecked(parent)){
            	safeUncheck(parent);
            }
		}
		repaint();
	}

	/**
	 * An uncheck method without the side effect of unchecking all
	 * children.
	 * Climbs the parent heirarchy and unchecks nodes until it finds one that isn't checked.
	 * @param node
	 */
	private void safeUncheck(final DefaultMutableTreeNode node){
		checkedNodes.remove(node);
		DefaultMutableTreeNode parent = (DefaultMutableTreeNode)node.getParent();
		if(parent != null && isNodeChecked(parent)){
				safeUncheck(parent);
		}
	}

	boolean isNodeChecked(final DefaultMutableTreeNode node){
		return checkedNodes.contains(node);
	}

    public TreePath[] getCheckedNodesPaths(){
        TreePath[] ret = new TreePath[checkedNodes.size()];
        for(int i = 0; i < checkedNodes.size(); i++){
            ret[i] = new TreePath(checkedNodes.get(i).getPath());
        }

        return ret;
    }

    public Vector<DefaultMutableTreeNode> getCheckedNodes(){
        return checkedNodes;
    }

    /**
     * Returns a shallow list of all the checked nodes.
     * i.e. a dfs is done through the tree and as soon as node is found
     * to be checked we dont go any deeper and move onto the next sibling.
     * This is usefull based on the fact that if a node is checked all of its
     * children must also be checked.
     *
     * @return
     */
    public ArrayList<DefaultMutableTreeNode> getCheckedNodesShallow(){
    	ArrayList<DefaultMutableTreeNode> nodes = new ArrayList<DefaultMutableTreeNode>();

    	getShallowHelper((DefaultMutableTreeNode)getModel().getRoot(), nodes);

    	return nodes;
    }

    /**
     * Recursive implementation of {@link #getCheckedNodesShallow()}
     * @param node
     * @param list
     */
    private void getShallowHelper(final DefaultMutableTreeNode node, final ArrayList<DefaultMutableTreeNode> list){
    	if(isNodeChecked(node)){
    		list.add(node);
    	}else{

    		for(int i = 0; i < node.getChildCount(); i++){
    			getShallowHelper((DefaultMutableTreeNode)node.getChildAt(i), list);

    		}

    	}
    }


    /**
     * Action handler for the change color menu items
     *
     */
    private class ColorChanger implements ActionListener {
        Color ourColor;
        public ColorChanger(final Color c) {
            ourColor = c;
        }
        public void actionPerformed(final ActionEvent e) {
            cellRenderer.setCheckedColor(ourColor);
            repaint();
        }
    };


	protected class CheckBoxTreeListener extends MouseAdapter{

        @Override
		public void mouseClicked(final MouseEvent e) {
            int chkBoxLocX = 0;
            int chkBoxLocY = cellRenderer.getCheckBoxYOffset();
            int chkBoxWidth = cellRenderer.getCheckBoxSize();
            int chkBoxHeight = cellRenderer.getCheckBoxSize();
			int x = e.getX();
			int y = e.getY();
			TreePath path = getPathForLocation(e.getX(), e.getY());
//			System.out.println("click at: " + e.getX() + ", " + e.getY());
			if(path != null){
				Object obj = path.getLastPathComponent();
//				System.out.println(obj);
				TreeUI tree = getUI();
				Rectangle rect = tree.getPathBounds(CheckBoxTree.this, path);

				if(x >= rect.getX() + chkBoxLocX - 2 && x <= rect.getX() + chkBoxLocX + chkBoxWidth + 2
						&& y >= rect.getY() + chkBoxLocY - 2 && y <= rect.getY() + chkBoxLocY + chkBoxHeight + 2){
//					System.out.println("Checked");
					checkNode((DefaultMutableTreeNode)obj);
				}else{
//					System.out.println("Not checked");
				}

			}else{
//				System.out.println("null path");
			}
		}
	}

    class ColorIcon implements Icon {

        private static final int DEFAULT_WIDTH = 10;
        private static final int DEFAULT_HEIGHT = 10;

        private Color iconColor;

        private final Dimension iconDimensions;

        public ColorIcon(final Color color) {
            this(color, DEFAULT_WIDTH, DEFAULT_HEIGHT);
        }

        public ColorIcon(final Color color, final int width, final int height){
            setIconColor(color);
            iconDimensions = new Dimension();
            setDimensions(width, height);
        }

        public void paintIcon(final Component c, final Graphics g, final int x, final int y) {
            Color startColor = g.getColor();
            g.setColor(getIconColor());
            g.fillOval(x + 1, y + 1, getIconWidth() - 2, getIconHeight() - 2);
            g.setColor(startColor);

        }

        public int getIconWidth() {
            return iconDimensions.width;
        }

        public int getIconHeight() {
            return iconDimensions.height;
        }

        public void setDimensions(final int width, final int height){
            iconDimensions.width = width;
            iconDimensions.height = height;
        }

        public Color getIconColor(){
            return iconColor;
        }

        public void setIconColor(final Color color){
            iconColor = color;
        }

    }

}
